package com.wxy.care.KnnUtils;

import org.springframework.beans.factory.annotation.Autowired;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;

public class KNN {

    //@Autowired
    private List<Instance> trainingSet; // 训练集
    private int k; // k值

    // 构造函数
    public KNN(int k) {
        this.k = k;
        trainingSet = new ArrayList<>();
    }

    // 添加实例到训练集
    public void addInstance(Instance instance) {
        trainingSet.add(instance);
    }

    // 预测实例的类别
    public String predict(Instance instance) {
        Distance[] distances = new Distance[trainingSet.size()];

        // 计算每个训练实例与待预测实例之间的距离
        for (int i = 0; i < trainingSet.size(); i++) {
            Distance distance = new Distance(trainingSet.get(i), instance);
            distances[i] = distance;
        }

        // 按照距离排序
        Arrays.sort(distances);

        // 取前k个最近的邻居
        HashMap<String, Integer> counts = new HashMap<>();
        for (int i = 0; i < k; i++) {
            Instance neighbor = distances[i].getInstance();
            String category = neighbor.getCategory();
            counts.put(category, counts.getOrDefault(category, 0) + 1);
        }

        // 找出出现次数最多的类别
        int maxCount = 0;
        String prediction = "";
        for (String category : counts.keySet()) {
            int count = counts.get(category);
            if (count > maxCount) {
                maxCount = count;
                prediction = category;
            }
        }

        return prediction;
    }

    // 计算实例之间的距离
    private static class Distance implements Comparable<Distance> {
        private final Instance instance1;
        private final Instance instance2;
        private final double distance;

        public Distance(Instance instance1, Instance instance2) {
            this.instance1 = instance1;
            this.instance2 = instance2;
            this.distance = calculateDistance();
        }

        public Instance getInstance() {
            return instance1;
        }

        private double calculateDistance() {
            double sum = 0.0;
            for (int i = 0; i < instance1.getFeatures().length; i++) {
                double diff = instance1.getFeatures()[i] - instance2.getFeatures()[i];
                sum += diff * diff;
            }
            return Math.sqrt(sum);
        }

        @Override
        public int compareTo(Distance other) {
            return Double.compare(distance, other.distance);
        }
    }
}
